import os
import cv2
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing import image

def load_model(model_path='garbage_classifier_model.keras'):
    """加载预训练模型"""
    if not os.path.exists(model_path):
        print(f"错误：模型文件 {model_path} 不存在")
        print("请确保：")
        print("1. 模型文件路径正确")
        print("2. 已经运行过训练脚本")
        return None
    
    try:
        model = tf.keras.models.load_model(model_path)
        print(f"成功加载模型: {model_path}")
        model.summary()  # 打印模型结构
        return model
    except Exception as e:
        print(f"加载模型时出错: {e}")
        return None

def capture_with_preview(model):
    """带预览的图像捕获和预测"""
    cap = cv2.VideoCapture(1)
    if not cap.isOpened():
        print("无法打开摄像头")
        return
    
    # 创建预览窗口
    cv2.namedWindow("垃圾识别预览", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("垃圾识别预览", 800, 600)
    
    print("\n摄像头预览已启动")
    print("按空格键捕获图像并进行分类")
    print("按ESC键退出")
    
    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法获取摄像头画面")
            break
        
        # 显示实时预览
        preview_frame = frame.copy()
        cv2.putText(preview_frame, "Press SPACE to capture", (10, 30),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
        cv2.putText(preview_frame, "Press ESC to quit", (10, 60),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        cv2.imshow("垃圾识别预览", preview_frame)
        
        key = cv2.waitKey(1)
        if key == 27:  # ESC键
            print("用户终止操作")
            break
        elif key == 32:  # 空格键
            # 捕获当前帧
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            # 预处理图像
            img = image.array_to_img(rgb_frame)
            img = img.resize((128, 128))  # 统一使用128x128尺寸
            img_array = image.img_to_array(img)
            img_array = img_array / 255.0  # 归一化到[0,1]
            img_array = np.expand_dims(img_array, axis=0)
            
            # 预测
            classes = ['cardboard', 'metal', 'plastic']
            pred_probs = model.predict(img_array)
            pred_label = np.argmax(pred_probs[0])
            confidence = np.max(pred_probs[0])
            
            # 显示结果
            print("\n" + "="*50)
            print("预测结果:")
            print(f"检测到: {classes[pred_label]}, 置信度: {confidence:.4f}")
            print("各类别概率:")
            for cls, prob in zip(classes, pred_probs[0]):
                print(f"  {cls}: {prob:.4f}")
            
            # 根据置信度阈值判断
            if confidence > 0.45:
                print(f"\n确定为垃圾: {classes[pred_label]}")
                # 这里可以添加机械臂控制代码
            else:
                print("\n未检测到有效垃圾或置信度过低")
            print("="*50 + "\n")
            
            # 在预览窗口显示结果
            result_frame = frame.copy()
            cv2.putText(result_frame, f"Class: {classes[pred_label]}", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
            cv2.putText(result_frame, f"Confidence: {confidence:.2f}", (10, 60),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
            cv2.imshow("垃圾识别预览", result_frame)
            cv2.waitKey(2000)  # 显示结果2秒
    
    cap.release()
    cv2.destroyAllWindows()

def main():
    # 加载模型
    model = load_model()
    if model is None:
        return
    
    # 启动带预览的捕获和预测
    capture_with_preview(model)

if __name__ == "__main__":
    main()